import argparse

import torch.backends.cudnn as cudnn

from models.experimental import *
from utils.datasets import *
from utils.utils import *

source = ['0', '0']


def detect():
    weights = 'params/helmet_head_person_s.pt'
    conf_thres = 0.8
    iou_thres = 0.5
    imgsz = 640

    # Initialize
    device = torch_utils.select_device()
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
    if half:
        model.half()  # to FP16

    # Set Dataloader
    cudnn.benchmark = True  # set True to speed up constant image size inference
    # dataset = LoadStreams('rtsp://admin:bocom123456@10.20.40.205:554', img_size=imgsz)
    dataset = LoadStreams(source, img_size=imgsz)
    # print("datasets:{}".format(dataset))

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]

    # Run inference
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        pred = model(img, augment=False)

        # Apply NMS
        # pred out: [tensor([[160.85034,  67.97704, 505.73584, 367.06537,   0.93149,   0.00000],
        #         [254.64270,  67.17798, 393.35028, 241.45377,   0.89275,   1.00000]])]
        # [left, top, right, bottom, conf, label_id]
        pred = non_max_suppression(pred[0], conf_thres, iou_thres, agnostic=True)

        # Process detections
        result = {}
        for i, det in enumerate(pred):  # detections per image
            result[i] = []
            p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
            if det is not None and len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                # Write results
                for *xyxy, conf, cls in det:
                    # print("xyxy:{}, conf:{}, cls:{}".format(xyxy, conf, cls))
                    label = '%s %.2f' % (names[int(cls)], conf)
                    # bbox: left, top, right, bottom
                    bbox = [int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])]
                    result[i].append({"bbox": bbox, "conf": round(float(conf), 3), "cls": names[int(cls)]})
                    plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
                    # Stream results
            print("i:{}, result:{}".format(i, result[i]))
            cv2.imshow(p, im0)
            if cv2.waitKey(1) == ord('q'):  # q to quit
                raise StopIteration


    print('Done. (%.3fs)' % (time.time() - t0))


if __name__ == '__main__':
    with torch.no_grad():
        detect()
